#!/usr/bin/python

import optparse
import os.path
import warnings

import numpy
import tables

import histogram as hist
from histogram.io import ptio

class BadCommandLineError(Exception):
    def __init__(self, msg, code):
        super(BadCommandLineError, self).__init__(msg)
        self.msg = msg
        self.retcode = code

    def __str__(self):
        return self.msg


def hmod_module_loader(fp, g_hmod):
    try:
        hmod = fp.getNodeAttr(g_hmod, 'module_path')
    except AttributeError:
        msg = "Found `hmod' node but node does not have `module_path' attr"
        raise AttributeError(msg)

    hc = hist.HistogramCollection()

    hmod_type = type(hmod)
    if hmod_type != type("hi!") and \
            hmod_type != type(numpy.string_('hi')):
        hc.module_name = hmod[0]
    else:
        hc.module_name = str(hmod)

    try:
        module = __import__(hc.module_name)
    except ImportError, e:
        msg = "HistogramCollection has module_name `{}', \n"
        msg += "but module could not be imported"
        warnings.warn(msg.format(hc.module_name))
        return None

    try:
        hcolls = module.hcolls
        hnames = module.hnames
    except AttributeError:
        try:
            hcolls = module.histdef.hcolls        
            hnames = module.histdef.hnames
        except AttributeError:
            msg = "Module imported but unable to locate histogram defs"
            warnings.warn(msg)
            return None

        
    for node in fp.iterNodes(fp.root, 'Group'):
        if node == g_hmod:
            continue

        k = node._v_name

        bins = node.bins.read()
        edgedict = {}
        for attr in "low_edges low_edges_x low_edges_y low_edges_z".split():
            try:
                edgedict[attr] = getattr(node, attr).read()
            except tables.NoSuchNodeError:
                pass
        
        if len(edgedict) == 0:
            msg = "Encountered key `{}' with no edge attrs"
            raise AttributeError(msg)

        if len(edgedict) == 1:
            xed = edgedict.values()[0]

            h = hist.Histogram1D(k, nbins=len(bins), 
                            arange=(xed.min(), xed.max()),
                            xlabel=k,
                            xunits='', bins=bins)
        
        elif len(edgedict) > 1:
            xed = edgedict['low_edges_x']
            yed = edgedict['low_edges_y']

            if len(edgedict) == 2:
                h = hist.Histogram2D(k, nbins=bins.shape, 
                                arange=((xed.min(), xed.max()),
                                        (yed.min(), yed.max())),
                                xlabel='x', ylabel='y',
                                bins=bins)
            else:
                zed = edgedict['low_edges_z']
                if zed.shape[0] == bins.shape[2]:
                    zed = numpy.linspace(zed[0], zed[-1], bins.shape[2]+1)

                h = hist.Histogram3D(k, nbins=bins.shape, 
                                arange=((xed.min(), xed.max()),
                                        (yed.min(), yed.max()),
                                        (zed.min(), zed.max())),
                                xlabel='x', ylabel='y', zlabel='z',
                                bins=bins)
                    
        hc.add(h)

    
    ptio.update_hist_attrs(hc)

    return hc


def hatt_module_loader(node, fp):
    msg = "Loading histograms from file with `hatt' attribute " + \
        "not yet supported!"

    raise NotImplementedError(msg)

old_hmod_keys = ['hmod', '__hatt__']
module_loaders = {'hmod': hmod_module_loader,
                  '__hatt__': hatt_module_loader,
                  }


def main():
    '''
    Usage: %prog [opts] file.hpt
    
    Read `file.hpt' (produced with an older version of the python-histogram
    library) and output those histograms in the current format.
    '''

    op = optparse.OptionParser(usage=main.__doc__)

    op.add_option("-n", "--name", dest='names', action='append', default=[],
                  help='Update histogram NAME only (can be repeated)',
                  metavar='NAME')    
    op.add_option("-o", "--output", dest='output', default=None, 
                  help='Leave output in PATH', metavar='PATH')

    opts, args = op.parse_args()

    if len(args) < 1:
        msg = "Please supply one `hpt' file to update"
        raise BadCommandLineError(msg, -1)

    path = args[0]

    if not tables.isPyTablesFile(path):
        msg = "Path `{0}' does not point to a PyTables file".format(path)
        raise IOError(msg)

    hc = None
    with tables.openFile(path) as fpin:
        for name in old_hmod_keys:
            try:
                hmod = fpin.getNode(fpin.root, name)
            except tables.NoSuchNodeError:
                pass
            else:
                hc = module_loaders[name](fpin, hmod)

    if opts.output is None:
        opts.output = os.path.basename(path).rsplit('.', 1)[0]
        opts.output += '.updated.hpt'
        print "Output left in `{}'".format(opts.output)

    hist.dump(hc, opts.output)

    return 0


if __name__ == '__main__':
    try:
        import sys
        sys.exit(main())
    except BadCommandLineError, e:
        import os
        sys.stderr.write(str(e) + "\n")
        sys.stderr.write("See `{0} -h' for more.\n".format(
                os.path.basename(sys.argv[0])))
        sys.stderr.flush()
        sys.exit(e.retcode)



